#include "solver.h"

/*** Constructor:
 * initialize mesh, DMDA, ksp solver, calls laplacian_matrix, allocate space and everything else
 * ***/
#undef __FUNCT__
#define __FUNCT__ "Solver"
Solver::Solver(integer Nx_, integer Ny_, bool adaptive_ts_, real lambda_, real domain_[2][2], real T_, std::string model_, MPI_Comm local_comm_, int local_color_)
{
  init(Nx_, Ny_, adaptive_ts_, lambda_, domain_, T_, model_, local_comm_, local_color_);
}

/*** Constructor:
 * from a options class
 * ***/
#undef __FUNCT__
#define __FUNCT__ "Solver"
Solver::Solver(Options *opt, MPI_Comm local_comm_, int local_color_)
{
  init(opt->Nx, opt->Ny, opt->adaptive, opt->lambda, opt->domain, opt->T, opt->model, local_comm_, local_color_);
  provide_initial_data();
  set_name(opt->name);
  set_snapshots(opt->take_snapshot, opt->num_snapshot, opt->times_snapshot);
}

/*** init:
 * initialize the class called by contructor
 * ***/
#undef __FUNCT__
#define __FUNCT__ "init"
PetscErrorCode Solver::init(integer Nx_, integer Ny_, bool adaptive_ts_, real lambda_, real domain_[2][2], real T_, std::string model_, MPI_Comm local_comm_, int local_color_)
{
  Nx = Nx_;
  Ny = Ny_;
  adaptive_ts = adaptive_ts_;
  lambda = lambda_;
  T = T_;
  model = model_;
  local_comm = local_comm_;
  local_color = local_color_;
  
  left = domain_[0][0];
  right = domain_[0][1];
  bottom = domain_[1][0];
  top = domain_[1][1];
  
  // Setup mesh
  dx = abs( right - left ) / Nx;
  dy = abs( top - bottom ) / Ny;
  domain[0][0] = left + dx/2;
  domain[0][1] = right - dx/2;
  domain[1][0] = bottom + dy/2;
  domain[1][1] = top - dy/2;
  area = abs( right - left ) * abs( top - bottom );
  
  step_counter = 0;
  take_snapshot = false;
  
  // Return this rank and number of processes
  //ierr = MPI_Comm_size(local_comm, &comm_size); CHKERRQ(ierr); // NO PARALLEL SIM
  //ierr = MPI_Comm_rank(local_comm, &loc_rank); CHKERRQ(ierr); // NO PARALLEL SIM
  ierr = MPI_Comm_size(local_comm, &comm_size); CHKERRQ(ierr);
  ierr = MPI_Comm_rank(local_comm, &loc_rank); CHKERRQ(ierr);
  
  // Create DMDA and setup
  ierr = DMDACreate2d(local_comm, DMDA_BOUNDARY_PERIODIC, DMDA_BOUNDARY_PERIODIC, DMDA_STENCIL_STAR,
                      Nx, Ny, PETSC_DECIDE,PETSC_DECIDE, 2, 2, NULL, NULL, &da); CHKERRQ(ierr);
  
  ierr = DMDASetUniformCoordinates(da, domain[0][0], domain[0][1], domain[1][0], domain[1][1], 0.0, 0.0); CHKERRQ(ierr);
  
  // Setup field names
  ierr = DMDASetFieldName(da,0,"x"); CHKERRQ(ierr);
  ierr = DMDASetFieldName(da,1,"y"); CHKERRQ(ierr);
  
  // Update local grid info
  ierr = DMDAGetCorners(da, &corn_x, &corn_y, 0, &span_x, &span_y, 0); CHKERRQ(ierr);
  ierr = DMDAGetGhostCorners(da, &ghost_corn_x, &ghost_corn_y, 0, &ghost_span_x, &ghost_span_y, 0); CHKERRQ(ierr);
  
  // Add mesh info
  integer               trash;
  DMDABoundaryType      trash_;
  DMDAStencilType       trash__;
  ierr = DMDAGetInfo(da, &trash, &trash, &trash, &trash, &proc_x, &proc_y, &trash, &trash, &trash, &trash_, &trash_, &trash_, &trash__); CHKERRQ(ierr);
  
  // Save DMDA info
  const integer *lx, *ly, *lz;
  vec_corn_x = new integer[comm_size];
  vec_corn_y = new integer[comm_size];
  vec_span_x = new integer[comm_size];
  vec_span_y = new integer[comm_size];
  cumu_areas = new integer[comm_size];
  DMDAGetOwnershipRanges(da, &lx, &ly, &lz);
  for( int i = 0; i < comm_size; i++ ) {
    integer c = i % proc_x;
    integer r = i / proc_x;
    
    vec_span_x[i] = lx[c];
    vec_span_y[i] = ly[r];
    vec_corn_x[i] = ( c!=0 ? vec_corn_x[i-1] + lx[c-1] : 0 );
    vec_corn_y[i] = ( r!=0 ? vec_corn_y[i-proc_x] + ly[r-1] : 0 );
    cumu_areas[i] = ( i!=0 ? cumu_areas[i-1] + vec_span_x[i-1] * vec_span_y[i-1] : 0 );
  }
  
#if VERBOSE
  PetscPrintf(local_comm,   "[INFO] Grid info: %dx%d mesh splitted as %dx%d between %d processors...\n",
              Nx, Ny, proc_x, proc_y, comm_size);
  
  PetscPrintf(local_comm,   "       DMDA has %d dof and a star stencil of size %d.\n", dof, stencil);
  
  for( int i = 0; i < comm_size; i++ ) {
    PetscPrintf(local_comm, "       Processor %d of %d get grid from corner [%d, %d], spanning [%d, %d] (not including ghosts), "
                                  "with a cumulative area of %d.\n",
                i, comm_size, vec_corn_x[i], vec_corn_y[i], vec_span_x[i], vec_span_y[i], cumu_areas[i]);
  }
#endif

#ifdef VERBOSE
  PetscPrintf(PETSC_COMM_SELF, "[INFO] (Process %d of %d, group %d) Starting memory allocation...\n", loc_rank, comm_size, local_color);
#endif

  // Create vectors
  ierr = DMCreateGlobalVector(da, &U_global); CHKERRQ(ierr);
  ierr = DMCreateLocalVector(da, &U_local); CHKERRQ(ierr);

  ierr = VecDuplicate(U_local, &G_local); CHKERRQ(ierr);
  ierr = VecDuplicate(U_local, &P_local); CHKERRQ(ierr);
  ierr = VecDuplicate(U_local, &U_x_local); CHKERRQ(ierr);
  ierr = VecDuplicate(U_local, &U_y_local); CHKERRQ(ierr);
  ierr = VecDuplicate(U_local, &U_no_x_local); CHKERRQ(ierr);
  ierr = VecDuplicate(U_local, &U_no_y_local); CHKERRQ(ierr);
  ierr = VecDuplicate(U_local, &U_edg_L_local); CHKERRQ(ierr);
  ierr = VecDuplicate(U_local, &U_edg_R_local); CHKERRQ(ierr);
  
  ierr = VecDuplicate(U_global, &G_global); CHKERRQ(ierr);
  ierr = VecDuplicate(U_global, &P_global); CHKERRQ(ierr);
  ierr = VecDuplicate(U_global, &U_x_global); CHKERRQ(ierr);
  ierr = VecDuplicate(U_global, &U_y_global); CHKERRQ(ierr);
  ierr = VecDuplicate(U_global, &U_no_x_global); CHKERRQ(ierr);
  ierr = VecDuplicate(U_global, &U_no_y_global); CHKERRQ(ierr);
  ierr = VecDuplicate(U_global, &U_edg_L_global); CHKERRQ(ierr);
  ierr = VecDuplicate(U_global, &U_edg_R_global); CHKERRQ(ierr);

  // Setup dummy vectors
  ierr = VecCreate(local_comm, &phi); CHKERRQ(ierr);
  ierr = VecSetSizes(phi, span_x * span_y, Nx * Ny); CHKERRQ(ierr);
  ierr = VecSetFromOptions(phi); CHKERRQ(ierr);
  ierr = VecSetUp(phi); CHKERRQ(ierr);

  ierr = VecCreate(local_comm, &Div_U); CHKERRQ(ierr);
  ierr = VecSetSizes(Div_U, span_x * span_y, Nx * Ny); CHKERRQ(ierr);
  ierr = VecSetFromOptions(Div_U); CHKERRQ(ierr);
  ierr = VecSetUp(Div_U); CHKERRQ(ierr);

  // Update grid ranges for linear vectors
  integer       start_, end_;
  ierr = VecGetOwnershipRange(Div_U, &start_, &end_); CHKERRQ(ierr);
  ierr = VecGetOwnershipRange(phi, &istart, &iend); CHKERRQ(ierr);
  
#ifdef VERBOSE // Print array info
  PetscPrintf(PETSC_COMM_SELF, "[INFO] (Process %d of %d, group %d) Local ownership range is: [%d, %d].\n",
              loc_rank, comm_size, local_color, istart, iend);
#endif
  
  // Create laplacian matrix
  ierr = laplacian_matrix(); CHKERRQ(ierr);
  
  // Create ksp from A
  ierr = KSPCreate(local_comm, &ksp); CHKERRQ(ierr);
  ierr = KSPSetOperators(ksp, A, A, SAME_PRECONDITIONER); CHKERRQ(ierr);
  ierr = KSPSetFromOptions(ksp); CHKERRQ(ierr);

#ifdef VERBOSE // Check sizes and print array info
 if( start_ != istart || end_ != iend ) {
    PetscPrintf(PETSC_COMM_SELF, "[ERROR] (Process %d of %d, group %d) FATAL ERROR: vector ownership do not coincide: [%d, %d] VS [%d, %d].\n",
                loc_rank, comm_size, local_color, istart, iend, start_, end_);
  }
  if( iend - istart != span_x * span_y ) {
    PetscPrintf(PETSC_COMM_SELF, "[ERROR] (Process %d of %d, group %d) FATAL ERROR: vector and DMDA split size ip do not coincide: %d VS %d.\n",
                loc_rank, comm_size, local_color, iend - istart, span_x * span_y);
  }
#endif

  return 0;
}

/*** Destructor:
 * delete what is remaining
 * ***/
#undef __FUNCT__
#define __FUNCT__ "~Solver"
Solver::~Solver()
{
#ifdef VERBOSE
  PetscPrintf(PETSC_COMM_SELF, "[INFO] Clearing workspace...\n");
#endif
  
  delete[] vec_corn_x;
  delete[] vec_corn_y;
  delete[] vec_span_x;
  delete[] vec_span_y;
  delete[] cumu_areas;
}

/*** deallocate:
 * delete everything, allocated memory, matrices, ksp, da and vectors.
 * ***/
#undef __FUNCT__
#define __FUNCT__ "deallocate"
PetscErrorCode Solver::deallocate() {
  
#ifdef VERBOSE
  PetscPrintf(PETSC_COMM_SELF, "[INFO] Deallocating useless memory...\n");
#endif
  
  // Finalize matrices
  ierr = KSPDestroy(&ksp); CHKERRQ(ierr);
  ierr = MatDestroy(&A); CHKERRQ(ierr);
  
  // Finalize vectors
  ierr = VecDestroy(&U_local); CHKERRQ(ierr);
  ierr = VecDestroy(&G_local); CHKERRQ(ierr);
  ierr = VecDestroy(&P_local); CHKERRQ(ierr);
  ierr = VecDestroy(&U_x_local); CHKERRQ(ierr);
  ierr = VecDestroy(&U_y_local); CHKERRQ(ierr);
  ierr = VecDestroy(&U_no_x_local); CHKERRQ(ierr);
  ierr = VecDestroy(&U_no_y_local); CHKERRQ(ierr);
  ierr = VecDestroy(&U_edg_L_local); CHKERRQ(ierr);
  ierr = VecDestroy(&U_edg_R_local); CHKERRQ(ierr);
  
  ierr = VecDestroy(&U_global); CHKERRQ(ierr);
  ierr = VecDestroy(&G_global); CHKERRQ(ierr);
  ierr = VecDestroy(&P_global); CHKERRQ(ierr);
  ierr = VecDestroy(&U_x_global); CHKERRQ(ierr);
  ierr = VecDestroy(&U_y_global); CHKERRQ(ierr);
  ierr = VecDestroy(&U_no_x_global); CHKERRQ(ierr);
  ierr = VecDestroy(&U_no_y_global); CHKERRQ(ierr);
  ierr = VecDestroy(&U_edg_L_global); CHKERRQ(ierr);
  ierr = VecDestroy(&U_edg_R_global); CHKERRQ(ierr);
  
  ierr = VecDestroy(&phi); CHKERRQ(ierr);
  ierr = VecDestroy(&Div_U); CHKERRQ(ierr);
  
  return 0;
}

/*** set_snapshots:
 * setup snapshot times and perform checking
 * ***/
#undef __FUNCT__
#define __FUNCT__ "set_snapshots"
void Solver::set_snapshots(bool take, integer num, real *times) {
  take_snapshot = take;
  num_snapshot = num;
  
  for(int i = 0; i < num; i++) {
    snapshot_times[i] = times[i];
    
#if VERBOSE
    if( i > 0 && snapshot_times[i] <= snapshot_times[i-1] ) PetscPrintf(local_comm, "[WARNING] Invalid snapshots at %f (s < s+1)...\n", snapshot_times[i]);
    if( snapshot_times[i] <= 0. || snapshot_times[i]>= T ) PetscPrintf(local_comm, "[WARNING] Invalid snapshots at %f (0 < s < T)...\n", snapshot_times[i]);
#endif
  }
}

/*** Provide_initial_data
 * set function pointers to correct functions to compute inital data later
 * ***/
#undef __FUNCT__
#define __FUNCT__ "provide_initial_data"
void Solver::provide_initial_data()
{
  initialize_data();
}

/*** Initialize_data:
 * inizialize vector U and G with function handles provided by u_0 etc...
 * ***/
#undef __FUNCT__
#define __FUNCT__ "initialize_data"
PetscErrorCode Solver::initialize_data()
{
    scalar        x, y;
    Field         **U, **G;

#if VERBOSE
    PetscPrintf(local_comm, "[INFO] (Group %d) Setting up initial data...\n", local_color);
#endif
    
    if( model == "shear-layer" ) {
      initial_data = new ContShear;
    } else if( model == "disc-shear-layer" ) {
      initial_data  = new DiscShear;
    } else if( model == "vortex-box" ) {
      initial_data  = new Vortex;
    } else if( model == "steady" ) {
      initial_data  = new Steady;
    }
#ifdef USE_RANDOM
    else if( model == "random-shear" ) {
      initial_data  = new RngShear;
    }
    else if( model == "random-interface" ) {
       initial_data  = new RngInterface;
    }
#endif
    
    ierr = DMDAVecGetArray(da, U_global, &U); CHKERRQ(ierr);
    ierr = DMDAVecGetArray(da, G_global, &G); CHKERRQ(ierr);
    for ( int j = corn_y; j < corn_y + span_y; j++ )
    {
      for ( int i = corn_x; i < corn_x + span_x; i++ )
      {
        
        get_coords(y, x, j, i);
#if  DEBUG
//       PetscPrintf(local_comm, "(%d x %d) --> [%f, %f] --|| %f, %f, %f, %f\n",
//                               j, i, y, x, u_0(x,y), v_0(x,y), grad_p_x_0(x,y), grad_p_y_0(x,y));
#endif
        if( model == "shear_layer" ) {
          U[j][i].x = static_cast<ContShear*>(initial_data)->u_0(x,y);
          U[j][i].y = static_cast<ContShear*>(initial_data)->v_0(x,y);
          G[j][i].x = static_cast<ContShear*>(initial_data)->g_p_x_0(x,y);
          G[j][i].y = static_cast<ContShear*>(initial_data)->g_p_y_0(x,y);
        } else if( model == "disc-shear-layer" ) {
          U[j][i].x = static_cast<DiscShear*>(initial_data)->u_0(x,y);
          U[j][i].y = static_cast<DiscShear*>(initial_data)->v_0(x,y);
          G[j][i].x = static_cast<DiscShear*>(initial_data)->g_p_x_0(x,y);
          G[j][i].y = static_cast<DiscShear*>(initial_data)->g_p_y_0(x,y);
        } else if( model == "vortex-box" ) {
          U[j][i].x = static_cast<Vortex*>(initial_data)->u_0(x,y);
          U[j][i].y = static_cast<Vortex*>(initial_data)->v_0(x,y);
          G[j][i].x = static_cast<Vortex*>(initial_data)->g_p_x_0(x,y);
          G[j][i].y = static_cast<Vortex*>(initial_data)->g_p_y_0(x,y);
        } else if( model == "steady" ) {
          U[j][i].x = static_cast<Steady*>(initial_data)->u_0(x,y);
          U[j][i].y = static_cast<Steady*>(initial_data)->v_0(x,y);
          G[j][i].x = static_cast<Steady*>(initial_data)->g_p_x_0(x,y);
          G[j][i].y = static_cast<Steady*>(initial_data)->g_p_y_0(x,y);
        }
#ifdef  USE_RANDOM
        else if( model == "random-shear" ) {
          U[j][i].x = static_cast<RngShear*>(initial_data)->u_0(x,y);
          U[j][i].y = static_cast<RngShear*>(initial_data)->v_0(x,y);
          G[j][i].x = static_cast<RngShear*>(initial_data)->g_p_x_0(x,y);
          G[j][i].y = static_cast<RngShear*>(initial_data)->g_p_y_0(x,y);
        }
        else if( model == "random-interface" ) {
          U[j][i].x = static_cast<RngInterface*>(initial_data)->u_0(x,y);
          U[j][i].y = static_cast<RngInterface*>(initial_data)->v_0(x,y);
          G[j][i].x = static_cast<RngInterface*>(initial_data)->g_p_x_0(x,y);
          G[j][i].y = static_cast<RngInterface*>(initial_data)->g_p_y_0(x,y);
        }
#endif
      }
    }
    ierr = DMDAVecRestoreArray(da, U_global, &U); CHKERRQ(ierr);
    ierr = DMDAVecRestoreArray(da, G_global, &G); CHKERRQ(ierr);
    
    ierr = DMGlobalToLocalBegin(da, U_global, INSERT_VALUES, U_local); CHKERRQ(ierr);
    ierr = DMGlobalToLocalEnd(da, U_global, INSERT_VALUES, U_local); CHKERRQ(ierr);
    ierr = DMGlobalToLocalBegin(da, G_global, INSERT_VALUES, G_local); CHKERRQ(ierr);
    ierr = DMGlobalToLocalEnd(da, G_global, INSERT_VALUES, G_local); CHKERRQ(ierr);
    
    return 0;
}

/*** Laplacian_matrix:
 * allocate and build the matrix A for the projection method
 * ***/
#undef __FUNCT__
#define __FUNCT__ "laplacian_matrix"
PetscErrorCode Solver::laplacian_matrix() {
  scalar         v, v_x, v_y;
  integer        C, L, R, B, T;
  
  ierr = MatCreate(local_comm, &A); CHKERRQ(ierr);
  ierr = MatSetSizes(A, span_x*span_y, span_x*span_y, Nx * Ny, Nx * Ny); CHKERRQ(ierr);
  ierr = MatSetFromOptions(A); CHKERRQ(ierr);
  ierr = MatMPIAIJSetPreallocation(A, 5, NULL, 3, NULL); CHKERRQ(ierr);
  ierr = MatSeqAIJSetPreallocation(A, 5, NULL); CHKERRQ(ierr);
  ierr = MatSetUp(A); CHKERRQ(ierr);
  
#if VERBOSE
  PetscPrintf(PETSC_COMM_SELF, "[INFO] (Process %d of %d, group %d) Assembly %dx%d Laplacian matrix... Local: %d -> %d.\n",
              loc_rank, comm_size, local_color, Nx*Ny, Nx*Ny, istart, iend);
#endif
  
  for ( C = istart; C < iend; C++ )
  {
    get_nhbrds(L, R, B, T, C);
    v_x = 1 / (dx * dx);
    v_y = 1 / (dy * dy);
    v = -2.0 * ( 1 / (dx * dx) + 1 / (dy * dy) );
    
    ierr = MatSetValues(A, 1, &C, 1, &L, &v_x, INSERT_VALUES); CHKERRQ(ierr);
    ierr = MatSetValues(A, 1, &C, 1, &R, &v_x, INSERT_VALUES); CHKERRQ(ierr);
    ierr = MatSetValues(A, 1, &C, 1, &B, &v_y, INSERT_VALUES); CHKERRQ(ierr);
    ierr = MatSetValues(A, 1, &C, 1, &T, &v_y, INSERT_VALUES); CHKERRQ(ierr);
    ierr = MatSetValues(A, 1, &C, 1, &C, &v, INSERT_VALUES); CHKERRQ(ierr);
  }
  
  ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  
  return 0;
}

/*** get_rank:
 * return rank of processor owning global index j,i
 * ***/
#undef __FUNCT__
#define __FUNCT__ "get_rank"
integer Solver::get_rank(const integer & j, const integer & i) {
  int col, row;
  for( int I = 0 ; I < proc_x; I++ ) {
    if( vec_corn_x[I] <= i ) col = I;
  }
  for( int J = 0 ; J < proc_y; J++) {
    if( vec_corn_y[J*proc_x+col] <= j ) row = J;
  }
  
  //PetscPrintf(PETSC_COMM_SELF, "[[ %d, %d, %d, %d ]] corn y\n", vec_corn_y[0], vec_corn_y[1], vec_corn_y[2], vec_corn_y[3]);
  //PetscPrintf(PETSC_COMM_SELF, "[[ %d, %d, %d, %d ]] corn x\n", vec_corn_x[0], vec_corn_x[1], vec_corn_x[2], vec_corn_x[3]);
  //PetscPrintf(PETSC_COMM_SELF, "[[ %d, %d ]] is in proc %d, %d\n", j, i, row, col);
  
  return col + row * proc_x;
}

/*** global_to_inline:
 * convert global index to local inline index
 * ***/
#undef __FUNCT__
#define __FUNCT__ "global_to_inline"
void Solver::global_to_inline(integer & I, const integer & j, const integer & i) {
  
  I = istart + (i - corn_x) + (j - corn_y) * span_x;
  
#ifdef VERBOSE
  if( istart > I || I >= iend ) PetscPrintf(PETSC_COMM_SELF, "[WARNING] (Process %d of %d, group %d) Index %d out of local ownership range [%d, %d].\n",
                                            loc_rank, comm_size, local_color, I, istart, iend);
  if( i < ghost_corn_x || i >= ghost_corn_x + ghost_span_x ) PetscPrintf(PETSC_COMM_SELF, "[WARNING] (Process %d of %d, group %d) Index %d out of local grid ghosted range [%d, %d].\n",
                                                                         loc_rank, comm_size, local_color, j, ghost_corn_x, ghost_corn_x + ghost_span_x);
  if( j < ghost_corn_y || j >= ghost_corn_y + ghost_span_y ) PetscPrintf(PETSC_COMM_SELF, "[WARNING] (Process %d of %d, group %d) Index %d out of local grid ghosted range [%d, %d].\n",
                                                                         loc_rank, comm_size, local_color, j, ghost_corn_y, ghost_corn_y + ghost_span_y);
#endif
}

/*** inline_to_global:
 * convert inline index to global index (j, i)
 * ***/
#undef __FUNCT__
#define __FUNCT__ "inline_to_global"
void Solver::inline_to_global(integer & j, integer & i, const integer & I) {
  integer rel_pos = (I - istart);
  j = corn_y + rel_pos / span_x;
  i = corn_x + rel_pos % span_x;
  
#ifdef VERBOSE
  if( istart > I || I >= iend ) PetscPrintf(PETSC_COMM_SELF, "[WARNING] (Process %d of %d, group %d) Index %d out of local ownership range [%d, %d].\n",
                                            loc_rank, comm_size, local_color, I, istart, iend);
  if( i < ghost_corn_x || i >= ghost_corn_x + ghost_span_x ) PetscPrintf(PETSC_COMM_SELF, "[WARNING] (Process %d of %d, group %d) Index %d out of local grid ghosted range [%d, %d].\n",
                                                                        loc_rank, comm_size, local_color, j, ghost_corn_x, ghost_corn_x + ghost_span_x);
  if( j < ghost_corn_y || j >= ghost_corn_y + ghost_span_y ) PetscPrintf(PETSC_COMM_SELF, "[WARNING] (Process %d of %d, group %d) Index %d out of local grid ghosted range [%d, %d].\n",
                                                                         loc_rank, comm_size, local_color, j, ghost_corn_y, ghost_corn_y + ghost_span_y);
#endif
}

/*** get_nhbrds
 * return neiborhoods of inline coord C in inline coords
 * ***/
#undef __FUNCT__
#define __FUNCT__ "get_nhbrds"
void Solver::get_nhbrds(integer & L, integer & R, integer & B, integer & T, const integer & C) {
  integer       j, i;
  inline_to_global(j, i, C);
  
  integer       r1 = get_rank(j, mod( i - 1, Nx));
  integer       r2 = get_rank(j, mod( i + 1, Nx));
  integer       r3 = get_rank(mod( j - 1, Ny) , i);
  integer       r4 = get_rank(mod( j + 1, Ny) , i);
  
  L = mod( cumu_areas[r1] + (j - vec_corn_y[r1]) * vec_span_x[r1] + (mod( i - 1, Nx) - vec_corn_x[r1]), Nx * Ny);
  R = mod( cumu_areas[r2] + (j - vec_corn_y[r2]) * vec_span_x[r2] + (mod( i + 1, Nx) - vec_corn_x[r2]), Nx * Ny);
  B = mod( cumu_areas[r3] + (mod( j - 1, Ny) - vec_corn_y[r3]) * vec_span_x[r3] + (i - vec_corn_x[r3]), Nx * Ny);
  T = mod( cumu_areas[r4] + (mod( j + 1, Ny) - vec_corn_y[r4]) * vec_span_x[r4] + (i - vec_corn_x[r4]), Nx * Ny);
  
  //PetscPrintf(PETSC_COMM_SELF, "## %d ** %d, %d **---> L = %d (%d, %d) (p. %d), R = %d (%d, %d) (p. %d), B = %d (%d, %d) (p. %d), T = %d (%d, %d) (p. %d)\n", C, j, i, L,j, mod( i - 1, Nx), r1, R,j, mod( i + 1, Nx), r2, B, mod( j - 1, Ny) , i, r3, T, mod( j + 1, Ny) , i, r4);
}

/*** get_coords
 * return x,y in function of global coords j,i
 * ***/
#undef __FUNCT__
#define __FUNCT__ "get_coords"
void Solver::get_coords(scalar & y, scalar & x, const integer & j, const integer & i) {
  x = domain[0][0] + i * dx;
  y = domain[1][0] + j * dy;
}

/*** solve:
 * perform actual solution including timestepping initialization etc...
 * ***/
#undef __FUNCT__
#define __FUNCT__ "solve"
PetscErrorCode Solver::solve() {
  
  real                  perc;
  real                  v1, v2, v1_cpu, v2_cpu;
  
  last_snapshot = 0;
  
#if VERBOSE
  PetscPrintf(local_comm, "[INFO] Starting new solution of the system up to time %.2f.\n", T);
#endif
  
  ierr = PetscTime(&v1); CHKERRQ(ierr);
  ierr = PetscGetCPUTime(&v1_cpu); CHKERRQ(ierr);
  
  times.push_back(0);
  VecNorm(U_global, NORM_2, &l2);
  energy.push_back(l2 * l2 * dx * dy / area);
  
  for( t = 0; t < T; )
  {
    VecNorm(U_global, NORM_INFINITY, &sup);
    adaptive_ts ? dt = min(dx, dy) / sup * lambda : dt = min(dx, dy) * lambda;
    //if( t + dt >= T ) dt = T - t;
    
    do_snapshot = false;
    if( take_snapshot && last_snapshot < num_snapshot && t <= snapshot_times[last_snapshot] && t + dt >= snapshot_times[last_snapshot] ) {
      //dt = snapshot_times[last_snapshot] - t;
      last_snapshot++;
      do_snapshot = true;
    }
    
    MPI_Bcast(&dt, 1, MPI_DOUBLE, 0, local_comm);
    MPI_Bcast(&do_snapshot, 1, MPI_INTEGER, 0, local_comm);
    
    perc = t / T * 100.;
    
#if VERBOSE
    PetscPrintf(local_comm, "[INFO] ---------------------- GROUP %d: NEW TIME STEP...  ----------------------\n", local_color);
    PetscPrintf(local_comm, "[INFO] [%3.2f%%] Step %d at time t = %f of T = %f: initiating new time step of size dt = %f (sup(U) = %f)..\n",
                perc, step_counter, t, T, dt, sup);
#endif
    
    ierr = step(); CHKERRQ(ierr);
    t += dt;
    
    times.push_back(t);
    VecNorm(U_global, NORM_2, &l2);
    energy.push_back(l2 * l2 * dx * dy / area);
    
    if( do_snapshot == true ) {
#if VERBOSE
      PetscPrintf(local_comm, "[INFO] Taking snapshot as promised...\n");
#endif
      ierr = create_to_zero(); CHKERRQ(ierr);
      //save_solution(VELOCITY, CURL);
      save_solution(VELOCITY, X_COMP);
      save_solution(VELOCITY, Y_COMP);
    }
    
    step_counter++;
  }
  
#if VERBOSE
  VecNorm(U_global, NORM_INFINITY, &sup);
  VecNorm(U_global, NORM_2, &l2);
  l2 = l2 * sqrt(dx * dy) / area;
  PetscPrintf(local_comm, "[INFO] Solution completed. Norms are: Linf = %f, L2 = %f...\n",
              sup, l2);
#endif
  
  ierr = DMLocalToGlobalBegin(da, U_local, INSERT_VALUES, U_global); CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(da, U_local, INSERT_VALUES, U_global); CHKERRQ(ierr);
  
  ierr = PetscTime(&v2); CHKERRQ(ierr);
  ierr = PetscGetCPUTime(&v2_cpu); CHKERRQ(ierr);
  
  elapsed = v2 - v1;
  elapsed_cpu = v2_cpu - v1_cpu;
  
#if VERBOSE
  PetscPrintf(local_comm, "[INFO] It took %d steps and time %f (CPU time %f)...\n", step_counter, elapsed, elapsed_cpu);
  PetscPrintf(local_comm, "[INFO] Saved %d solution snapshots...\n", last_snapshot);
#endif
  
  ierr = create_to_zero(); CHKERRQ(ierr);
  ierr = deallocate(); CHKERRQ(ierr);
  
  //save_solution(VELOCITY, CURL);
  save_solution(VELOCITY, X_COMP);
  save_solution(VELOCITY, Y_COMP);
  //save_energy();
  //save_log();
  
  return 0;
}

/*** step:
 * perform a single step for the method
 * ***/
#undef __FUNCT__
#define __FUNCT__ "step"
PetscErrorCode Solver::step() {
  
#if VERBOSE
  PetscPrintf(local_comm, "[INFO] Beginning new time step...\n", loc_rank, comm_size);
#endif
  
  VecCopy(U_local, P_local); // Copy U_local to P_local, we work on P_local
  
  ierr = DMLocalToGlobalBegin(da, P_local, INSERT_VALUES, P_global); CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(da, P_local, INSERT_VALUES, P_global); CHKERRQ(ierr);
  
  ierr = convection(); CHKERRQ(ierr); // P_local now contains the convection of U_local
  
  VecAXPBY(P_local, -1, -1, G_local); // P_local containst the argument to the projection
  
  // Update bounareis of the convection
  // TODO
  ierr = DMLocalToGlobalBegin(da, P_local, INSERT_VALUES, P_global); CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(da, P_local, INSERT_VALUES, P_global); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(da, P_global, INSERT_VALUES, P_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, P_global, INSERT_VALUES, P_local); CHKERRQ(ierr);
  
  ierr = projection(); CHKERRQ(ierr); // Project P_local, get U_x (P(arg) and U_y Q(arg))
  
  ierr = DMLocalToGlobalBegin(da, U_local, INSERT_VALUES, U_global); CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(da, U_local, INSERT_VALUES, U_global); CHKERRQ(ierr);
  
  VecAXPY(U_global, dt, U_x_global); // U now contains the new step of U
  VecAXPY(G_global, 1, U_y_global); // G now contains the new step of U
  
  ierr = DMGlobalToLocalBegin(da, U_global, INSERT_VALUES, U_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, U_global, INSERT_VALUES, U_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(da, G_global, INSERT_VALUES, G_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, G_global, INSERT_VALUES, G_local); CHKERRQ(ierr);
  
  return 0;
}

/*** convection:
 * Read from P_local and write to P_global the convection of P
 * ***/
#undef __FUNCT__
#define __FUNCT__ "convection"
PetscErrorCode Solver::convection() {
  
  Field **U, **G, **U_x, **U_y, **U_no_x, **U_no_y, **U_edg_L, **U_edg_R;
  scalar L, R, C;
  
  //// First limit/*ed central difference, will be limited (MC) (without dx and dy division)
  ierr = DMDAVecGetArray(da, P_local, &U); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_x_global, &U_x); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_y_global, &U_y); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      L = U[j][i].x - U[j][i-1].x;
      R = U[j][i+1].x - U[j][i].x;
      C = ( U[j][i+1].x - U[j][i-1].x ) * 0.5;
      U_x[j][i].x = MC(L, R, C);
      
      L = U[j][i].y - U[j][i-1].y;
      R = U[j][i+1].y - U[j][i].y;
      C = ( U[j][i+1].y - U[j][i-1].y ) * 0.5;
      U_x[j][i].y = MC(L, R, C);
      
      L = U[j][i].x - U[j-1][i].x;
      R = U[j+1][i].x - U[j][i].x;
      C = ( U[j+1][i].x - U[j-1][i].x ) * 0.5;
      U_y[j][i].x = MC(L, R, C);
      
      L = U[j][i].y - U[j-1][i].y;
      R = U[j+1][i].y - U[j][i].y;
      C = ( U[j+1][i].y - U[j-1][i].y ) * 0.5;
      U_y[j][i].y = MC(L, R, C);
    }
  }
  ierr = DMDAVecRestoreArray(da, U_x_global, &U_x); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_y_global, &U_y); CHKERRQ(ierr);
  
  ierr = DMGlobalToLocalBegin(da, U_x_global, INSERT_VALUES, U_x_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, U_x_global, INSERT_VALUES, U_x_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(da, U_y_global, INSERT_VALUES, U_y_local); CHKERRQ(ierr);
  
  ierr = DMDAVecGetArray(da, U_x_local, &U_x); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_no_x_global, &U_no_x); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      //// Normal derivative and upwinding
      // x
      if (U[j][i].x >= 0) {
        C = 0.5 * (1 - dt / dx * U[j][i].x);
        U_no_x[j][i].x = ( U[j][i].x - U[j][i-1].x + C * ( U_x[j][i].x - U_x[j][i-1].x ) ) / dx;
        U_no_x[j][i].y = ( U[j][i].y - U[j][i-1].y + C * ( U_x[j][i].y - U_x[j][i-1].y ) ) / dx;
      } else {
        C = 0.5 * (1 + dt / dx * U[j][i].x);
        U_no_x[j][i].x = ( U[j][i+1].x - U[j][i].x - C * ( U_x[j][i+1].x - U_x[j][i].x ) ) / dx;
        U_no_x[j][i].y = ( U[j][i+1].y - U[j][i].y - C * ( U_x[j][i+1].y - U_x[j][i].y ) ) / dx;
      }
    }
  }
  ierr = DMDAVecRestoreArray(da, U_x_local, &U_x); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_no_x_global, &U_no_x); CHKERRQ(ierr);
  
  
  ierr = DMGlobalToLocalEnd(da, U_y_global, INSERT_VALUES, U_y_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(da, U_no_x_global, INSERT_VALUES, U_no_x_local); CHKERRQ(ierr);
  
  ierr = DMDAVecGetArray(da, U_y_local, &U_y); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_no_y_global, &U_no_y); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      // y
      if (U[j][i].y >= 0) {
        C = 0.5 * (1 - dt / dy * U[j][i].y);
        U_no_y[j][i].x = ( U[j][i].x - U[j-1][i].x + C * ( U_y[j][i].x - U_y[j-1][i].x ) ) / dy;
        U_no_y[j][i].y = ( U[j][i].y - U[j-1][i].y + C * ( U_y[j][i].y - U_y[j-1][i].y ) ) / dy;
      } else {
        C = 0.5 * (1 + dt / dy * U[j][i].y);
        U_no_y[j][i].x = ( U[j+1][i].x - U[j][i].x - C * ( U_y[j+1][i].x - U_y[j][i].x ) ) / dy;
        U_no_y[j][i].y = ( U[j+1][i].y - U[j][i].y - C * ( U_y[j+1][i].y - U_y[j][i].y ) ) / dy;
      }
    }
  }
  ierr = DMDAVecRestoreArray(da, U_y_local, &U_y); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_no_y_global, &U_no_y); CHKERRQ(ierr);
  
  ierr = DMGlobalToLocalEnd(da, U_no_x_global, INSERT_VALUES, U_no_x_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(da, U_no_y_global, INSERT_VALUES, U_no_y_local); CHKERRQ(ierr);
  
  ierr = DMDAVecGetArray(da, G_local, &G); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_no_x_local, &U_no_x); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_y_local, &U_y); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_edg_L_global, &U_edg_L); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_edg_R_global, &U_edg_R); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      C =  U[j][i].x - 0.5 * dt * ( U[j][i].y * U_y[j][i].x / dx + U[j][i].x * U_no_x[j][i].x + G[j][i].x );
      U_edg_L[j][i].x = C + 0.5 * U_y[j][i].x;
      U_edg_R[j][i].x = C - 0.5 * U_y[j][i].x;
      C =  U[j][i].y - 0.5 * dt * ( U[j][i].y * U_y[j][i].y / dy + U[j][i].x * U_no_x[j][i].y + G[j][i].y );
      U_edg_L[j][i].y = C + 0.5 * U_y[j][i].y;
      U_edg_R[j][i].y = C - 0.5 * U_y[j][i].y;
    }
  }
  ierr = DMDAVecRestoreArray(da, U_no_x_local, &U_no_x); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_y_local, &U_y); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_edg_L_global, &U_edg_L); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_edg_R_global, &U_edg_R); CHKERRQ(ierr);
  
  ierr = DMGlobalToLocalEnd(da, U_no_y_global, INSERT_VALUES, U_no_y_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(da, U_edg_L_global, INSERT_VALUES, U_edg_L_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, U_edg_L_global, INSERT_VALUES, U_edg_L_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(da, U_edg_R_global, INSERT_VALUES, U_edg_R_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, U_edg_R_global, INSERT_VALUES, U_edg_R_local); CHKERRQ(ierr);
  
  ierr = DMDAVecGetArray(da, U_edg_L_local, &U_edg_L); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_edg_R_local, &U_edg_R); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_y_global, &U_y); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      if( U_edg_L[j][i].y >= 0 && U_edg_L[j][i].y + U_edg_R[j+1][i].y >= 0 ) {
        U_y[j][i].y = U_edg_L[j][i].y;
      } else if( U_edg_L[j][i].y < 0 && U_edg_R[j+1][i].y > 0 ) {
        U_y[j][i].y = 0;
      } else {
        U_y[j][i].y = U_edg_R[j+1][i].y;
      }
      
      if( U_y[j][i].y > 0 ) {
        U_y[j][i].x = U_edg_L[j][i].x;
      } else if( U_y[j][i].y < 0 ) {
        U_y[j][i].x = U_edg_R[j+1][i].x;
      } else {
        U_y[j][i].x = 0.5 * ( U_edg_L[j][i].x +  U_edg_R[j+1][i].x );
      }
    }
  }
  ierr = DMDAVecRestoreArray(da, U_edg_L_local, &U_edg_L); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_edg_R_local, &U_edg_R); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_y_global, &U_y); CHKERRQ(ierr);
  
  ierr = DMGlobalToLocalBegin(da, U_y_global, INSERT_VALUES, U_y_local); CHKERRQ(ierr);
  
  ierr = DMDAVecGetArray(da, U_no_y_local, &U_no_y); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_x_local, &U_x); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_edg_L_global, &U_edg_L); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_edg_R_global, &U_edg_R); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      C =  U[j][i].x - 0.5 * dt * ( U[j][i].x * U_x[j][i].x / dx + U[j][i].y * U_no_y[j][i].x + G[j][i].x );
      U_edg_L[j][i].x = C + 0.5 * U_x[j][i].x;
      U_edg_R[j][i].x = C - 0.5 * U_x[j][i].x;
      C =  U[j][i].y - 0.5 * dt * ( U[j][i].x * U_x[j][i].y / dy + U[j][i].y * U_no_y[j][i].y + G[j][i].y );
      U_edg_L[j][i].y = C + 0.5 * U_x[j][i].y;
      U_edg_R[j][i].y = C - 0.5 * U_x[j][i].y;
    }
  }
  ierr = DMDAVecRestoreArray(da, G_local, &G); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_no_y_local, &U_no_y); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_x_local, &U_x); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_edg_L_global, &U_edg_L); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_edg_R_global, &U_edg_R); CHKERRQ(ierr);
  
  ierr = DMGlobalToLocalEnd(da, U_y_global, INSERT_VALUES, U_y_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(da, U_edg_L_global, INSERT_VALUES, U_edg_L_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, U_edg_L_global, INSERT_VALUES, U_edg_L_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(da, U_edg_R_global, INSERT_VALUES, U_edg_R_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, U_edg_R_global, INSERT_VALUES, U_edg_R_local); CHKERRQ(ierr);
  
  ierr = DMDAVecGetArray(da, U_edg_L_local, &U_edg_L); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_edg_R_local, &U_edg_R); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_x_global, &U_x); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      if( U_edg_L[j][i].x >= 0 && U_edg_L[j][i].x + U_edg_R[j][i+1].x >= 0 ) {
        U_x[j][i].x = U_edg_L[j][i].x;
      } else if( U_edg_L[j][i].x < 0 && U_edg_R[j][i+1].x > 0 ) {
        U_x[j][i].x = 0;
      } else {
        U_x[j][i].x = U_edg_R[j][i+1].x;
      }
      
      if( U_x[j][i].x > 0) {
        U_x[j][i].y = U_edg_L[j][i].y;
      } else if( U_x[j][i].x < 0) {
        U_x[j][i].y = U_edg_R[j][i+1].y;
      } else {
        U_x[j][i].y = 0.5 * ( U_edg_L[j][i].y + U_edg_R[j][i+1].y );
      }
    }
  }
  ierr = DMDAVecRestoreArray(da, U_edg_L_local, &U_edg_L); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_edg_R_local, &U_edg_R); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_x_global, &U_x); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, P_local, &U); CHKERRQ(ierr);
  
  ierr = DMGlobalToLocalBegin(da, U_x_global, INSERT_VALUES, U_x_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, U_x_global, INSERT_VALUES, U_x_local); CHKERRQ(ierr);
  
  ierr = DMDAVecGetArray(da, U_x_local, &U_x); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_y_local, &U_y); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, P_global, &U); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      L = (U_x[j][i-1].x + U_x[j][i].x) / dx;
      R = (U_y[j-1][i].y + U_y[j][i].y) / dy;
      U[j][i].x = -0.5 * (L * (U_x[j][i-1].x - U_x[j][i].x) + R * (U_y[j-1][i].x - U_y[j][i].x));
      U[j][i].y = -0.5 * (L * (U_x[j][i-1].y - U_x[j][i].y) + R * (U_y[j-1][i].y - U_y[j][i].y));
    }
  }
  ierr = DMDAVecRestoreArray(da, U_x_local, &U_x); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_y_local, &U_y); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, P_global, &U); CHKERRQ(ierr);
  
  ierr = DMGlobalToLocalBegin(da, P_global, INSERT_VALUES, P_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, P_global, INSERT_VALUES, P_local); CHKERRQ(ierr);
  
  return 0;
}

/*** projection:
 * Compute the projection of P and writes it onto U_x and U_y
 * ***/
#undef __FUNCT__
#define __FUNCT__ "projection"
PetscErrorCode Solver::projection() {
  Field      **P, **proj, **grad_phi;
  
#if VERBOSE
  PetscPrintf(local_comm, "[INFO] (Group %d) Computing divergence (RHS of system).\n", local_color, loc_rank, comm_size);
#endif
  ierr = divergence(); CHKERRQ(ierr);
  
#if VERBOSE
  PetscPrintf(local_comm, "[INFO] (Group %d) Solving linear system (may take some time).\n", local_color, loc_rank, comm_size);
#endif
  ierr = KSPSolve(ksp, Div_U, phi); CHKERRQ(ierr);
#if VERBOSE
  PetscPrintf(local_comm, "[INFO] (Group %d) System solved...\n", local_color, loc_rank, comm_size);
#endif
  
#if VERBOSE
  PetscPrintf(local_comm, "[INFO] (Group %d) Computing gradient (of solution)...\n", local_color, loc_rank, comm_size);
#endif
  ierr = gradient(); CHKERRQ(ierr);
  
  ierr = DMDAVecGetArray(da, U_y_local, &grad_phi); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, P_local, &P); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_x_global, &proj); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      
      proj[j][i].x = P[j][i].x - grad_phi[j][i].x;
      proj[j][i].y = P[j][i].y - grad_phi[j][i].y;
    }
  }
  ierr = DMDAVecRestoreArray(da, U_y_local, &grad_phi); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, P_local, &P); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_x_global, &proj); CHKERRQ(ierr);
  
  ierr = DMLocalToGlobalBegin(da, U_y_local, INSERT_VALUES, U_y_global); CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(da, U_y_local, INSERT_VALUES, U_y_global); CHKERRQ(ierr);
  
  ierr = DMGlobalToLocalBegin(da, U_x_global, INSERT_VALUES, U_x_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, U_x_global, INSERT_VALUES, U_x_local); CHKERRQ(ierr);
  
  return 0;
}

/*** gradient:
 * read from phi and writes the gradient to U_y, using U_x temporarly
 * ***/
#undef __FUNCT__
#define __FUNCT__ "gradient"
PetscErrorCode Solver::gradient() {
  
  Field      **grad_phi, **grad_phi_temp;
  integer       I;
  scalar        val;
  
  ierr = DMDAVecGetArray(da, U_no_x_global, &grad_phi_temp); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      global_to_inline(I, j, i);
      ierr = VecGetValues(phi, 1, &I, &val); CHKERRQ(ierr);
      grad_phi_temp[j][i].x = val;
    }
  }
  ierr = DMDAVecRestoreArray(da, U_no_x_global, &grad_phi_temp); CHKERRQ(ierr);
  
  ierr = DMGlobalToLocalBegin(da, U_no_x_global, INSERT_VALUES, U_no_x_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, U_no_x_global, INSERT_VALUES, U_no_x_local); CHKERRQ(ierr);
  
  ierr = DMDAVecGetArray(da, U_no_x_local, &grad_phi_temp); CHKERRQ(ierr);
  ierr = DMDAVecGetArray(da, U_y_global, &grad_phi); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ )
  {
    for ( int i = corn_x; i < corn_x + span_x; i++ )
    {
      grad_phi[j][i].x = 0.5 * ( ( grad_phi_temp[j][i+1].x - grad_phi_temp[j][i-1].x ) / dx);
      grad_phi[j][i].y = 0.5 * ( ( grad_phi_temp[j+1][i].x - grad_phi_temp[j-1][i].x ) / dy);
    }
  }
  ierr = DMDAVecRestoreArray(da, U_no_x_local, &grad_phi_temp); CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(da, U_y_global, &grad_phi); CHKERRQ(ierr);
  
  ierr = DMGlobalToLocalBegin(da, U_y_global, INSERT_VALUES, U_y_local); CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(da, U_y_global, INSERT_VALUES, U_y_local); CHKERRQ(ierr);
  
  return 0;
}

/*** divergence:
 * compute divergence of P_local and write into Div_U
 * ***/
#undef __FUNCT__
#define __FUNCT__ "divergence"
PetscErrorCode Solver::divergence() {
  
  Field      **D;
  scalar        val;
  integer       I;
  
  ierr = DMDAVecGetArray(da, P_local, &D); CHKERRQ(ierr);
  for ( int j = corn_y; j < corn_y + span_y; j++ ) {
    for ( int i = corn_x; i < corn_x + span_x; i++ ) {
      global_to_inline(I, j, i);
      val = 0.5 * ( (D[j+1][i].y - D[j-1][i].y) / dy  + (D[j][i+1].x - D[j][i-1].x) / dx );
      ierr = VecSetValues(Div_U, 1, &I, &val, INSERT_VALUES);
    }
  }
  ierr = DMDAVecRestoreArray(da, P_local, &D); CHKERRQ(ierr);
  
  ierr = VecAssemblyBegin(Div_U); CHKERRQ(ierr);
  ierr = VecAssemblyEnd(Div_U); CHKERRQ(ierr);
  
  return 0;
}

/*** create_to_zero
 * collect data from all processes and build U_collect, G_collect
 * ***/
#undef __FUNCT__
#define __FUNCT__ "create_to_zero"
PetscErrorCode Solver::create_to_zero() {
    VecScatter            U_ctx, G_ctx;
    Vec                   U_natural, G_natural;
    
#if VERBOSE
    PetscPrintf(local_comm, "[INFO] (Group %d) Collecting data to process 0.\n", local_color);
#endif
    
    ierr = DMDACreateNaturalVector(da, &U_natural); CHKERRQ(ierr);
    ierr = DMDACreateNaturalVector(da, &G_natural); CHKERRQ(ierr);
    
    ierr = VecScatterCreateToZero(U_global, &U_ctx, &U_collect); CHKERRQ(ierr);
    ierr = VecScatterCreateToZero(G_global, &G_ctx, &G_collect); CHKERRQ(ierr);
    
    ierr = DMDAGlobalToNaturalBegin(da, U_global, INSERT_VALUES, U_natural); CHKERRQ(ierr);
    ierr = DMDAGlobalToNaturalEnd(da, U_global, INSERT_VALUES, U_natural); CHKERRQ(ierr);
    ierr = DMDAGlobalToNaturalBegin(da, G_global, INSERT_VALUES, G_natural); CHKERRQ(ierr);
    ierr = DMDAGlobalToNaturalEnd(da, G_global, INSERT_VALUES, G_natural); CHKERRQ(ierr);
    
    ierr = VecScatterBegin(U_ctx, U_natural, U_collect, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
    ierr = VecScatterEnd(U_ctx, U_natural, U_collect, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
    ierr = VecScatterBegin(G_ctx, G_natural, G_collect, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
    ierr = VecScatterEnd(G_ctx, G_natural, G_collect, INSERT_VALUES, SCATTER_FORWARD); CHKERRQ(ierr);
    
    ierr = VecScatterDestroy(&U_ctx); CHKERRQ(ierr);
    ierr = VecDestroy(&U_natural); CHKERRQ(ierr);
    ierr = VecScatterDestroy(&G_ctx); CHKERRQ(ierr);
    ierr = VecDestroy(&G_natural); CHKERRQ(ierr);
    
    return 0;
}

/*** save_solution
 * save vector to dat and export with gnuplot (optional)
 * ***/
#undef __FUNCT__
#define __FUNCT__ "save_solution"
PetscErrorCode Solver::save_solution(QuantityType quantity, DataType data, PlotType plot)
{
  if( loc_rank != 0 ) return 0;
  
  Vec                   *vector;
  scalar                *V;
  integer               I, I_p, I_m, J_p, J_m, S = 2 * Nx * Ny;
  
  FILE                  *f;
  char                  fname[PETSC_MAX_PATH_LEN], dat_fname[PETSC_MAX_PATH_LEN], pdf_fname[PETSC_MAX_PATH_LEN], command[PETSC_MAX_PATH_LEN];
  
  if( quantity == VELOCITY ) {
    vector  = &U_collect;
  } else if ( quantity == PRESSURE ) {
    vector  = &G_collect;
  }
  
  if( data == X_COMP ) {
    sprintf(fname, "%s%s_x-comp_%dx%d_%.2f_nproc_%d_sim_%d", PATH_TO_DATA, name, Nx, Ny, t, comm_size, local_color);
  } else if( data == Y_COMP ) {
    sprintf(fname, "%s%s_y-comp_%dx%d_%.2f_nproc_%d_sim_%d", PATH_TO_DATA, name, Nx, Ny, t, comm_size, local_color);
  } else if( data == CURL ) {
    sprintf(fname, "%s%s_curl_%dx%d_%.2f_nproc_%d_sim_%d", PATH_TO_DATA, name, Nx, Ny, t, comm_size, local_color);
  } else if( data == DIV ) {
    sprintf(fname, "%s%s_div_%dx%d_%.2f_nproc_%d_sim_%d", PATH_TO_DATA, name, Nx, Ny, t, comm_size, local_color);
  }
  strcpy(dat_fname, fname); strcat(dat_fname, ".dat");
  strcpy(pdf_fname, fname); strcat(pdf_fname, ".pdf");
   
#if VERBOSE
  PetscPrintf(PETSC_COMM_SELF, "[INFO] (Process %d of %d, group %d) Exporting solution data to \"%s\"...\n",
              loc_rank, comm_size, local_color, dat_fname);
#endif
  
  f = fopen(dat_fname, "w");
  VecGetArray(*vector, &V);
  for ( int j = 0; j < Ny; j++ )
  {
    for ( int i = 0; i < Nx; i++ )
    {
    if( data == X_COMP ) {
      I = (i + j * Nx) * 2;
      ierr = PetscFPrintf(PETSC_COMM_SELF, f, "%G ", V[I]);
    } else if ( data == Y_COMP ) {
      I = (i + j * Nx) * 2 + 1;
      ierr = PetscFPrintf(PETSC_COMM_SELF, f, "%G ", V[I]);
    } else if ( data == CURL ) {
      I = (i + j * Nx) * 2;
      I_p = I + 1 + 2;                    (I_p >= S) && (I_p -= S);
      I_m = I + 1 - 2;                    (I_m < 0)  && (I_m += S);
      J_p = I + 2 * Nx;      (J_p >= S) && (J_p -= S);
      J_m = I - 2 * Nx;      (J_m < 0)  && (J_m += S);
      ierr = PetscFPrintf(PETSC_COMM_SELF, f, "%G ", ( V[I_p] - V[I_m] ) / dy - ( V[J_p] - V[J_m] ) / dx );
    } else if ( data == DIV ) {
      I = (i + j * Nx) * 2;
      I_p = I + 2;                    (I_p >= S) && (I_p -= S);
      I_m = I - 2;                    (I_m < 0)  && (I_m += S);
      J_p = I + 1 + 2 * Nx;      (J_p >= S) && (J_p -= S);
      J_m = I + 1 - 2 * Nx;      (J_m < 0)  && (J_m += S);
      ierr = PetscFPrintf(PETSC_COMM_SELF, f, "%G ", ( V[I_p] - V[I_m] ) / dx + ( V[J_p] - V[J_m] ) / dy );
    }
  }
  ierr = PetscFPrintf(PETSC_COMM_SELF, f, "\n"); CHKERRQ(ierr);
  }
  VecRestoreArray(*vector, &V);
  fclose(f);
  
  if( plot != NONE ) {
    integer usepdf;
    if( plot == PDF) usepdf = 1;
    else usepdf = 0;
  
    sprintf(command, "gnuplot -e \"inname='%s';outname='%s';Nx=%d;Ny=%d;L=%f;R=%f;B=%f;T=%f;time=%f;usepdf=%d;group=%d\" './tools/solution.gnuplot'",
            dat_fname, pdf_fname, Nx, Ny, left,  right, bottom, top, t, usepdf, local_color);
#if VERBOSE
    PetscPrintf(PETSC_COMM_SELF, "[INFO] (Process %d of %d, group %d) Running \"%s\"...\n",
                loc_rank, comm_size, local_color, command);
#endif
    system(command);
  }
  
  return 0;
}

/*** save_energy
 * save energy to a file an plot (optional)
 * ***/
#undef __FUNCT__
#define __FUNCT__ "save_energy"
PetscErrorCode Solver::save_energy(PlotType plot)
{
  if( loc_rank != 0 ) return 0;
  
  FILE                  *f;
  char                  fname[PETSC_MAX_PATH_LEN], dat_fname[PETSC_MAX_PATH_LEN], pdf_fname[PETSC_MAX_PATH_LEN], command[PETSC_MAX_PATH_LEN];
  
  sprintf(fname, "%s%s_energy_%dx%d_%.2f_nproc_%d_sim_%d", PATH_TO_DATA, name, Nx, Ny, T, comm_size, local_color);
  
  strcpy(dat_fname, fname); strcat(dat_fname, ".dat");
  strcpy(pdf_fname, fname); strcat(pdf_fname, ".pdf");
  
  #if VERBOSE
  PetscPrintf(PETSC_COMM_SELF, "[INFO] (Process %d of %d, group %d) Exporting energy data to \"%s\"...\n",
              loc_rank, comm_size, local_color, dat_fname);
  #endif
  
  f = fopen(dat_fname, "w");
  
  for ( unsigned int i = 0; i < times.size(); i++ )
  {
    ierr = PetscFPrintf(PETSC_COMM_SELF, f, "%G %G\n", times.at(i), energy.at(i)); CHKERRQ(ierr);
  }
  
  fclose(f);
  
  if( plot != NONE ) {
    integer usepdf;
    if( plot == PDF) usepdf = 1;
    else usepdf = 0;
    
    sprintf(command, "gnuplot -e \"inname='%s';outname='%s';Nx=%d;Ny=%d;L=%f;R=%f;B=%f;T=%f;usepdf=%d;group=%d\" './tools/energy.gnuplot'",
            dat_fname, pdf_fname, Nx, Ny, domain[0][0],  domain[0][0], domain[0][0], domain[0][0], usepdf, local_color);
#if VERBOSE
    PetscPrintf(PETSC_COMM_SELF, "[INFO] (Process %d of %d, group %d) Running \"%s\"...\n",
                loc_rank, comm_size, local_color, command);
#endif
    system(command);
  }
  
  return 0;
}

/*** save_log
 * save log to a file
 * ***/
#undef __FUNCT__
#define __FUNCT__ "save_log"
PetscErrorCode Solver::save_log()
{
  if( loc_rank != 0 ) return 0;
  
  FILE                  *f;
  char                  fname[PETSC_MAX_PATH_LEN], txt_fname[PETSC_MAX_PATH_LEN];
  
  sprintf(fname, "%s%s_report_%dx%d_%.2f_nproc_%d_sim_%d", PATH_TO_DATA, name, Nx, Ny, T, comm_size, local_color);
  
  strcpy(txt_fname, fname); strcat(txt_fname, ".log");
  
#if VERBOSE
  PetscPrintf(PETSC_COMM_SELF, "[INFO] (Process %d of %d, group %d) Exporting log to \"%s\"...\n",
              loc_rank, comm_size, local_color, txt_fname);
#endif
  
  f = fopen(txt_fname, "w");
  ierr = PetscFPrintf(local_comm, f, "*** Report generated the date %d by process %d of %d, group %d. ***\n\n", 0,
                      loc_rank, comm_size, local_color); CHKERRQ(ierr);
  
  ierr = PetscFPrintf(local_comm, f, "[1] GENERAL INFORMATION\n"); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Name:                   %s\n", name); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Time stepping:          %s\n", adaptive_ts ==  true ? "adaptive" : "constant"); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Lambda:                 %.2f\n", lambda); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Final time:             %.2f\n", T); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Domain:                 [%.4f, %.4f] x [%.4f, %.4f]\n", left, right, bottom, top); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Grid size:              %d x %d\n", Nx, Ny); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "\n"); CHKERRQ(ierr);
  
  ierr = PetscFPrintf(local_comm, f, "[2] PROCESS INFORMATION\n"); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Number of processes:    %d\n", comm_size); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Processor grid:         %dx%d\n", proc_x, proc_y); CHKERRQ(ierr);
  for(int i = 0; i < comm_size; i++) {
    ierr = PetscFPrintf(local_comm, f, "     Processor %d:\n", i); CHKERRQ(ierr);
    ierr = PetscFPrintf(local_comm, f, "       Grid corner:          [%d, %d]\n", vec_corn_x[i], vec_corn_y[i]); CHKERRQ(ierr);
    ierr = PetscFPrintf(local_comm, f, "       Grid span:            [%d, %d]\n", vec_span_x[i], vec_span_y[i]); CHKERRQ(ierr);
    ierr = PetscFPrintf(local_comm, f, "       Grid area:            %d\n", vec_span_x[i] * vec_span_y[i]); CHKERRQ(ierr);
  }
  ierr = PetscFPrintf(local_comm, f, "\n"); CHKERRQ(ierr);
  
  ierr = PetscFPrintf(local_comm, f, "[3] SOLUTION INFORMATION\n"); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Supremum:               %f\n", sup); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     L2 norm:                %f\n", l2); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Elapsed time:           %f\n", elapsed); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "     Elapsed CPU time:       %f\n", elapsed_cpu); CHKERRQ(ierr);
  ierr = PetscFPrintf(local_comm, f, "\n"); CHKERRQ(ierr);
  
  fclose(f);
  
  return 0;
}
